package com.tlgen.orm.components.offHeap;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.util.Pool;
import com.tlgen.orm.config.SpeedOrmConfig;
import com.tlgen.orm.utils.SpringUtils;
import org.springframework.scheduling.annotation.Scheduled;
import sun.misc.Unsafe;

import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Pattern;

/**
 * 堆外存储管理器 - 提供高性能、低GC的对象存储能力
 */
public class OffHeapManager {
    private static final Unsafe UNSAFE;
    private static final ConcurrentHashMap<String, MemoryBlock> STORAGE = new ConcurrentHashMap<>();
    private static final AtomicLong TOTAL_MEMORY = new AtomicLong(0);

    // 使用配置的最大内存限制
    private static volatile long MAX_MEMORY;

    // 熔断器状态
    private static volatile boolean circuitBreakerOpen = false;
    private static volatile long circuitBreakerOpenTime = 0;
    private static final AtomicInteger totalOperations = new AtomicInteger(0);
    private static final AtomicInteger errorOperations = new AtomicInteger(0);

    // 访问频率跟踪
    private static final Map<String, AtomicInteger> accessFrequency = new HashMap<>();
    private static final Map<String, Long> lastAccessTime = new HashMap<>();

    // 内存池管理
    private static final MemoryPool smallMemoryPool = new MemoryPool();
    private static final MemoryPool mediumMemoryPool = new MemoryPool();
    private static final MemoryPool largeMemoryPool = new MemoryPool();

    // 使用IdentityHashMap来跟踪已访问的对象，防止循环引用
    private static final ThreadLocal<IdentityHashMap<Object, Object>> VISITED_CACHE =
            ThreadLocal.withInitial(IdentityHashMap::new);

    // 创建 Kryo 对象池
    private static final Pool<Kryo> kryoPool = new Pool<Kryo>(true, false, 8) {
        protected Kryo create() {
            Kryo kryo = new Kryo();
            kryo.setRegistrationRequired(false); // 不需要注册类
            kryo.setReferences(false); // 关闭引用跟踪提高性能
            return kryo;
        }
    };

    static {
        try {
            // 获取Unsafe实例
            Field theUnsafe = Unsafe.class.getDeclaredField("theUnsafe");
            theUnsafe.setAccessible(true);
            UNSAFE = (Unsafe) theUnsafe.get(null);

            // 注册JVM关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(OffHeapManager::cleanupAll));

            // 初始化默认最大内存
            MAX_MEMORY = Runtime.getRuntime().maxMemory() / 4;
        } catch (Exception e) {
            throw new RuntimeException("Failed to initialize Unsafe", e);
        }
    }

    /**
     * 初始化配置 - 需要在应用启动后调用一次
     */
    public static void initConfiguration() {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config.getOffHeap() != null) {
            // 设置最大内存
            SpeedOrmConfig.OffHeapConfig offheapConfig = config.getOffHeap();
            SpeedOrmConfig.OffHeapConfig.Memory memoryConfig = offheapConfig.getMemory();
            if (memoryConfig != null && memoryConfig.getMaxSize() != null) {
                MAX_MEMORY = memoryConfig.getMaxSize().toBytes();
            }

            // 初始化内存池大小
            if (memoryConfig != null && memoryConfig.getPools() != null) {
                smallMemoryPool.init(memoryConfig.getPools().getSmall().toBytes());
                mediumMemoryPool.init(memoryConfig.getPools().getMedium().toBytes());
                largeMemoryPool.init(memoryConfig.getPools().getLarge().toBytes());
            }
        }
    }

    /**
     * 检查是否需要迁移到堆外存储
     * @param key 对象键
     * @param obj 待评估对象
     * @return 是否应迁移到堆外存储
     */
    public static boolean shouldMigrateToOffHeap(String key, Object obj) {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config.getOffHeap() == null || !config.getOffHeap().isEnabled()) {
            return false;
        }

        SpeedOrmConfig.OffHeapConfig.Migration migrationConfig = config.getOffHeap().getMigration();
        if (migrationConfig == null) {
            return false;
        }

        // 1. 检查大小阈值
        long size = estimateObjectSize(obj);
        if (size > migrationConfig.getSizeThreshold()) {
            return true;
        }

        // 2. 检查访问频率阈值
        if (accessFrequency.containsKey(key)) {
            int accesses = accessFrequency.get(key).get();
            if (accesses > migrationConfig.getAccessThreshold()) {
                return true;
            }
        }

        return false;
    }

    // 任何对象只要估算大小超过 1024 字节就会被放入堆外存储
    private static long estimateObjectSize(Object obj) {
        IdentityHashMap<Object, Object> visited = VISITED_CACHE.get();
        try {
            visited.clear();
            return estimateObjectSizeHelper(obj, visited);
        } finally {
            visited.clear();
        }
    }

    private static long estimateObjectSizeHelper(Object obj, IdentityHashMap<Object, Object> visited) {
        if (obj == null) return 0;
        if (visited.containsKey(obj)) return 0; // 已访问过，避免循环引用

        visited.put(obj, null);

        long size = 0;

        // 1. 对象头开销（12-16字节）
        size += System.getProperty("java.vm.name").contains("64") ? 16 : 12;

        // 处理数组类型
        if (obj.getClass().isArray()) {
            Class<?> componentType = obj.getClass().getComponentType();
            int length = Array.getLength(obj);

            if (componentType.isPrimitive()) {
                // 原始类型数组
                if (componentType == long.class || componentType == double.class) {
                    size += 8L * length;
                } else if (componentType == int.class || componentType == float.class) {
                    size += 4L * length;
                } else if (componentType == short.class || componentType == char.class) {
                    size += 2L * length;
                } else {
                    size += length; // byte/boolean
                }
            } else {
                // 对象引用数组
                size += 4L * length; // 引用大小
                for (int i = 0; i < length; i++) {
                    Object element = Array.get(obj, i);
                    if (element != null) {
                        size += estimateObjectSizeHelper(element, visited);
                    }
                }
            }
            return size;
        }

        // 处理非数组对象
        for (Field field : obj.getClass().getDeclaredFields()) {
            field.setAccessible(true);
            Class<?> fieldType = field.getType();

            if (fieldType.isPrimitive()) {
                if (fieldType == long.class || fieldType == double.class) size += 8;
                else if (fieldType == int.class || fieldType == float.class) size += 4;
                else if (fieldType == short.class || fieldType == char.class) size += 2;
                else size += 1; // byte/boolean
            } else {
                try {
                    Object fieldValue = field.get(obj);
                    if (fieldValue != null) {
                        size += estimateObjectSizeHelper(fieldValue, visited);
                    } else {
                        size += 4; // 空引用大小
                    }
                } catch (IllegalAccessException e) {
                    size += 4; // 访问失败按最小引用大小计算
                }
            }
        }

        return size;
    }

    private static long estimateFieldSize(Object fieldValue) {
        if (fieldValue instanceof String) {
            // String对象大小 = 对象头 + char数组
            return 16 + ((String) fieldValue).length() * 2L;
        } else if (fieldValue instanceof Collection) {
            // 集合大小 = 集合对象 + 元素大小
            Collection<?> col = (Collection<?>) fieldValue;
            long size = 40; // 基础集合对象大小
            for (Object item : col) {
                size += estimateObjectSize(item);
            }
            return size;
        } else {
            // 其他对象递归估算
            return estimateObjectSize(fieldValue);
        }
    }


    private static byte[] serialize(Object obj) {
        Kryo kryo = kryoPool.obtain();
        try (Output output = new Output(1024, -1)) {
            kryo.writeClassAndObject(output, obj);
            return output.toBytes();
        } finally {
            kryoPool.free(kryo);
        }
    }

    private static <T> T deserialize(byte[] data, Class<T> type) {
        Kryo kryo = kryoPool.obtain();
        try (Input input = new Input(data)) {
            return type.cast(kryo.readClassAndObject(input));
        } finally {
            kryoPool.free(kryo);
        }
    }

    /**
     * 判断对象是否应存储在堆外
     * @param obj 待评估对象
     * @return 是否应使用堆外存储
     */
    public static boolean shouldStoreOffHeap(Object obj) {
        if (obj == null) return false;

        // 估算对象大小
        long size = estimateObjectSize(obj);

        // 使用配置的迁移阈值判断
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config.getOffHeap() != null && config.getOffHeap().isEnabled()) {
            SpeedOrmConfig.OffHeapConfig.Migration migration = config.getOffHeap().getMigration();
            if (migration != null) {
                return size > migration.getSizeThreshold();
            }
        }

        // 默认阈值检查 (1KB)
        return size > 1024;
    }

    /**
     * 跟踪对象访问频率
     * @param key 对象键
     */
    public static void trackAccess(String key) {
        long now = System.currentTimeMillis();

        // 每分钟重置一次计数器
        if (lastAccessTime.containsKey(key)) {
            long lastTime = lastAccessTime.get(key);
            if (now - lastTime > 60000) { // 超过1分钟
                accessFrequency.computeIfAbsent(key, k -> new AtomicInteger(0)).set(0);
                lastAccessTime.put(key, now);
            }
        } else {
            lastAccessTime.put(key, now);
            accessFrequency.computeIfAbsent(key, k -> new AtomicInteger(0));
        }

        // 增加访问计数
        accessFrequency.get(key).incrementAndGet();
    }

    /**
     * 在保存操作前处理堆外存储
     * @param entity 实体对象
     */
    public static void beforeSave(Object entity) {
        if (shouldStoreOffHeap(entity)) {
            String key = generateKey(entity);
            put(key, entity);
        }
    }

    /**
     * 在保存操作前处理可能的迁移
     * @param key 对象键
     * @param entity 实体对象
     */
    public static void handleMigration(String key, Object entity) {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config == null || config.getOffHeap() == null || !config.getOffHeap().isEnabled()) {
            return;
        }

        // 如果对象已在堆外，不处理
        if (STORAGE.containsKey(key)) {
            return;
        }

        // 检查是否需要迁移
        if (shouldMigrateToOffHeap(key, entity)) {
            put(key, entity);
        }
    }

    /**
     * 在查询操作后尝试从堆外获取
     * @param id 实体ID
     * @param type 实体类型
     * @return 堆外存储的实体或null
     */
    public static <T> T afterSelect(Object id, Class<T> type) {
        String key = generateKey(id, type);
        return get(key, type);
    }

    /**
     * 检查熔断器状态
     * @return 熔断器是否开启
     */
    private static boolean isCircuitBreakerOpen() {
        if (!circuitBreakerOpen) return false;

        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config == null || config.getOffHeap() == null) {
            return false;
        }

        SpeedOrmConfig.OffHeapConfig.CircuitBreaker circuitBreaker = config.getOffHeap().getCircuitBreaker();
        if (circuitBreaker == null) {
            return false;
        }

        long now = System.currentTimeMillis();
        if (now - circuitBreakerOpenTime > circuitBreaker.getRecoveryTimeout()) {
            // 超过恢复时间，关闭熔断器
            circuitBreakerOpen = false;
            return false;
        }

        return true;
    }

    /**
     * 记录操作错误
     */
    private static void recordError() {
        errorOperations.incrementAndGet();
        totalOperations.incrementAndGet();
        checkCircuitBreaker();
    }

    /**
     * 记录操作成功
     */
    private static void recordSuccess() {
        totalOperations.incrementAndGet();
    }

    /**
     * 检查是否需要打开熔断器
     */
    private static void checkCircuitBreaker() {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config.getOffHeap() == null) {
            return;
        }

        SpeedOrmConfig.OffHeapConfig.CircuitBreaker circuitBreaker = config.getOffHeap().getCircuitBreaker();
        if (circuitBreaker == null) {
            return;
        }

        int total = totalOperations.get();
        int errors = errorOperations.get();

        if (total > 0) {
            double errorRate = (double) errors * 100 / total;
            if (errorRate > circuitBreaker.getErrorThreshold()) {
                circuitBreakerOpen = true;
                circuitBreakerOpenTime = System.currentTimeMillis();
            }
        }
    }

    /**
     * 将对象存储到堆外内存
     * @param key 存储键
     * @param obj 存储对象
     */
    public static void put(String key, Object obj) {
        if (key == null || obj == null) return;

        // 检查熔断器
        if (isCircuitBreakerOpen()) {
            return;
        }

        // 增加操作计数
        try {
            // 序列化对象
            byte[] data = serialize(obj);

            // 检查内存限制
            if (TOTAL_MEMORY.get() + data.length > MAX_MEMORY) {
                if (!tryEvict(data.length)) {
                    throw new OutOfOffHeapMemoryException("Exceeded max off-heap memory");
                }
            }

            // 根据对象大小选择合适的内存池
            MemoryPool pool = getMemoryPool(data.length);
            long address;
            if (pool != null && pool.hasCapacity()) {
                address = pool.allocate(data.length);
            } else {
                // 如果池中没有足够空间，直接分配
                address = UNSAFE.allocateMemory(data.length);
            }

            // 复制数据到堆外
            UNSAFE.copyMemory(data, Unsafe.ARRAY_BYTE_BASE_OFFSET, null, address, data.length);

            // 创建内存块记录
            MemoryBlock block = new MemoryBlock(address, data.length, System.currentTimeMillis(), pool != null);

            // 存储到映射表
            STORAGE.put(key, block);
            TOTAL_MEMORY.addAndGet(data.length);

            // 记录成功
            recordSuccess();

        } catch (Exception e) {
            // 记录错误
            recordError();
            throw new OffHeapStorageException("Failed to store object off-heap", e);
        }
    }

    /**
     * 根据大小选择合适的内存池
     * @param size 数据大小
     * @return 最适合的内存池或null
     */
    private static MemoryPool getMemoryPool(int size) {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        if (config == null || config.getOffHeap() == null || config.getOffHeap().getMemory() == null) {
            return null;
        }

        if (size <= smallMemoryPool.getBlockSize()) {
            return smallMemoryPool;
        } else if (size <= mediumMemoryPool.getBlockSize()) {
            return mediumMemoryPool;
        } else if (size <= largeMemoryPool.getBlockSize()) {
            return largeMemoryPool;
        }

        return null;
    }

    /**
     * 从堆外内存获取对象
     * @param key 存储键
     * @param type 对象类型
     * @return 反序列化的对象
     */
    public static <T> T get(String key, Class<T> type) {
        if (key == null) return null;

        // 检查熔断器
        if (isCircuitBreakerOpen()) {
            return null;
        }

        // 增加访问跟踪
        trackAccess(key);

        MemoryBlock block = STORAGE.get(key);
        if (block == null) return null;

        try {
            // 从堆外内存读取数据
            byte[] data = new byte[block.size];
            UNSAFE.copyMemory(null, block.address, data, Unsafe.ARRAY_BYTE_BASE_OFFSET, block.size);

            // 反序列化对象
            T result = deserialize(data, type);

            // 记录成功
            recordSuccess();
            return result;
        } catch (Exception e) {
            // 记录错误
            recordError();

            // 移除损坏的数据
            STORAGE.remove(key);
            release(block);
            return null;
        }
    }

    /**
     * 释放堆外内存
     * @param key 存储键
     */
    public static void release(String key) {
        MemoryBlock block = STORAGE.remove(key);
        if (block != null) {
            release(block);
        }
    }

    /**
     * 获取当前堆外内存使用情况
     * @return 已使用字节数
     */
    public static long getMemoryUsage() {
        return TOTAL_MEMORY.get();
    }

    /**
     * 清理所有堆外内存
     */
    public static void cleanupAll() {
        STORAGE.keySet().forEach(OffHeapManager::release);
        STORAGE.clear();

        // 清理内存池
        smallMemoryPool.cleanup();
        mediumMemoryPool.cleanup();
        largeMemoryPool.cleanup();
    }

    // ========== 私有方法 ==========

    private static void release(MemoryBlock block) {
        if (block != null) {
            if (block.pooled) {
                // 如果是池中分配的内存，归还给内存池
                MemoryPool pool = getMemoryPool(block.size);
                if (pool != null) {
                    pool.release(block.address);
                } else {
                    UNSAFE.freeMemory(block.address);
                }
            } else {
                UNSAFE.freeMemory(block.address);
            }
            TOTAL_MEMORY.addAndGet(-block.size);
        }
    }

    private static boolean tryEvict(int requiredSpace) {
        // 使用LRU策略清除最旧的数据
        long oldestTimestamp = Long.MAX_VALUE;
        String oldestKey = null;

        for (String key : STORAGE.keySet()) {
            MemoryBlock block = STORAGE.get(key);
            if (block.timestamp < oldestTimestamp) {
                oldestTimestamp = block.timestamp;
                oldestKey = key;
            }
        }

        if (oldestKey != null) {
            release(oldestKey);
            return TOTAL_MEMORY.get() + requiredSpace <= MAX_MEMORY;
        }

        // 如果没有找到可清除的对象，尝试清除内存池中的空间
        if (largeMemoryPool.tryRelease(requiredSpace) ||
                mediumMemoryPool.tryRelease(requiredSpace) ||
                smallMemoryPool.tryRelease(requiredSpace)) {
            return true;
        }

        return false;
    }

    public static String generateKey(Object entity) {
        return entity.getClass().getName() + "@" + System.identityHashCode(entity);
    }

    public static String generateKey(Object id, Class<?> type) {
        return type.getName() + "@" + id.toString();
    }

    // ========== 内部类 ==========

    private static class MemoryBlock {
        final long address;
        final int size;
        final long timestamp;
        final boolean pooled; // 是否来自内存池

        MemoryBlock(long address, int size, long timestamp, boolean pooled) {
            this.address = address;
            this.size = size;
            this.timestamp = timestamp;
            this.pooled = pooled;
        }
    }

    /**
     * 内存池管理类
     */
    private static class MemoryPool {
        private volatile long poolAddress = 0;
        private volatile long poolSize = 0;
        private volatile long allocated = 0;
        private final Map<Long, Integer> allocatedBlocks = new HashMap<>(); // address -> size

        public void init(long size) {
            if (size <= 0) return;

            try {
                poolSize = size;
                poolAddress = UNSAFE.allocateMemory(size);
                allocated = 0;
                allocatedBlocks.clear();
            } catch (Exception e) {
                poolAddress = 0;
                poolSize = 0;
            }
        }

        public boolean hasCapacity() {
            return poolAddress != 0 && allocated < poolSize;
        }

        public boolean hasCapacity(long size) {
            return poolAddress != 0 && (allocated + size) <= poolSize;
        }

        public long getBlockSize() {
            return poolSize;
        }

        public long allocate(long size) {
            if (!hasCapacity(size)) {
                return 0;
            }

            long address = poolAddress + allocated;
            allocated += size;
            allocatedBlocks.put(address, (int) size);
            return address;
        }

        public void release(long address) {
            if (allocatedBlocks.containsKey(address)) {
                int size = allocatedBlocks.get(address);
                allocated -= size;
                allocatedBlocks.remove(address);

                // 清空内存
                UNSAFE.setMemory(address, size, (byte) 0);
            }
        }

        public boolean tryRelease(long requiredSpace) {
            if (!allocatedBlocks.isEmpty()) {
                // 找到第一个块释放
                Long address = allocatedBlocks.keySet().iterator().next();
                release(address);
                return requiredSpace <= (poolSize - allocated);
            }
            return false;
        }

        public void cleanup() {
            if (poolAddress != 0) {
                UNSAFE.freeMemory(poolAddress);
                poolAddress = 0;
                poolSize = 0;
                allocated = 0;
                allocatedBlocks.clear();
            }
        }
    }

    // 懒加载配置
    private static class Holder {
        static final SpeedOrmConfig SPEED_ORM_CONFIG_INSTANCE = SpringUtils.getBean(SpeedOrmConfig.class);
    }

    // 定期清理过期缓存
    @Scheduled(fixedRate = 30 * 60 * 1000) // 每30分钟
    public static void cleanExpiredOffHeapCache() {
        SpeedOrmConfig config = Holder.SPEED_ORM_CONFIG_INSTANCE;
        long expirationPeriod = 24 * 60 * 60 * 1000; // 默认24小时

        // 如果可以配置过期时间，可以在这里添加

        long threshold = System.currentTimeMillis() - expirationPeriod;

        List<String> keysToRemove = new ArrayList<>();
        STORAGE.forEach((key, block) -> {
            if (block.timestamp < threshold) {
                keysToRemove.add(key);
            }
        });

        keysToRemove.forEach(OffHeapManager::release);
    }

    // ======================================


    public static class OutOfOffHeapMemoryException extends RuntimeException {
        public OutOfOffHeapMemoryException(String message) {
            super(message);
        }
    }

    public static class OffHeapStorageException extends RuntimeException {
        public OffHeapStorageException(String message, Throwable cause) {
            super(message, cause);
        }
    }

    /**
     * 按实体类型清理缓存
     * @param entityClass 实体类
     * @return 清理的缓存项数量
     */
    public static int removeByEntityType(Class<?> entityClass) {
        String className = entityClass.getName();
        return removeByPattern("^" + Pattern.quote(className) + "@.*");
    }

    /**
     * 按正则表达式模式清理缓存
     * @param pattern 正则表达式模式
     * @return 清理的缓存项数量
     */
    public static int removeByPattern(String pattern) {
        Pattern regex = Pattern.compile(pattern);
        List<String> keysToRemove = new ArrayList<>();

        // 第一步：收集匹配的键
        STORAGE.keySet().forEach(key -> {
            if (regex.matcher(key).matches()) {
                keysToRemove.add(key);
            }
        });

        // 第二步：批量清理
        keysToRemove.forEach(OffHeapManager::release);
        return keysToRemove.size();
    }

    /**
     * 按实体类型和ID清理缓存
     * @param entityClass 实体类
     * @param id 实体ID
     * @return 是否成功清理
     */
    public static boolean removeByEntityId(Class<?> entityClass, Object id) {
        String key = generateKey(id, entityClass);
        if (STORAGE.containsKey(key)) {
            release(key);
            return true;
        }
        return false;
    }

    /**
     * 清理所有缓存
     * @return 清理的总内存大小（字节）
     */
    public static long clearAll() {
        long totalFreed = TOTAL_MEMORY.get();
        cleanupAll();
        return totalFreed;
    }

    /**
     * 获取缓存键列表（用于调试）
     * @return 所有缓存键的列表
     */
    public static List<String> getCacheKeys() {
        return new ArrayList<>(STORAGE.keySet());
    }

}
